//*********************************************************
//
//    Copyright (c) Microsoft. All rights reserved.
//    This code is licensed under the Apache License, Version 2.0.
//    THIS CODE IS PROVIDED *AS IS* WITHOUT WARRANTY OF
//    ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY
//    IMPLIED WARRANTIES OF FITNESS FOR A PARTICULAR
//    PURPOSE, MERCHANTABILITY, OR NON-INFRINGEMENT.
//
//********************************************************* 

using System;
using System.IO;
using System.Security.AccessControl;
using System.Security.Cryptography;
using System.Text;
using Microsoft.TridentWordAddIn.Common;
using Microsoft.Win32;
using SecurityAlgorithm;

namespace Microsoft.TridentWordAddIn.Presentation
{
    internal sealed class DataCryptor
    {
        private readonly string registryKeyName = "TridentWordAddIn_Key";
        private readonly string keyName = "KeyData";
        private readonly string iVName = "IVData";
        private readonly string registryKeyPath;

        #region Constructors

        /// <summary>
        /// Initializes a new instance of the <see cref="DataCryptor"/> class.
        /// </summary>
        public DataCryptor()
        {
            this.registryKeyPath = @"HKEY_CURRENT_USER\" + registryKeyName;
        }

        #endregion

        #region Enumeration
        /// <summary>
        /// Crypt Type
        /// </summary>
        private enum CryptType
        {
            /// <summary>
            /// Encryption
            /// </summary>
            Encrypt,
            /// <summary>
            /// Decryption
            /// </summary>
            Decrypt
        };

        #endregion

        #region Public Members

        /// <summary>
        /// Encrypts the and store data.
        /// </summary>
        /// <param name="sourceStream">The source stream.</param>
        /// <param name="destinationStream">The destination stream.</param>
        public void EncryptAndStoreData(MemoryStream sourceStream, Stream destinationStream)
        {
            ICryptoTransform cryptoTransform = GetCryptoTransform(CryptType.Encrypt);
            using (CryptoStream cryptoStream = new CryptoStream(destinationStream, cryptoTransform, CryptoStreamMode.Write))
            {
                StreamHelper.CopyStream(sourceStream, cryptoStream);
                cryptoStream.FlushFinalBlock();
            }
        }

        /// <summary>
        /// Gets the decrypted data.
        /// </summary>
        /// <param name="encryptedStream">The encrypted stream.</param>
        /// <returns></returns>
        public MemoryStream GetDecryptedData(Stream encryptedStream)
        {            
            ICryptoTransform cryptoTransform = GetCryptoTransform(CryptType.Decrypt);
            MemoryStream decryptedStream = new MemoryStream();

            using (CryptoStream cryptoStream = new CryptoStream(encryptedStream, cryptoTransform, CryptoStreamMode.Read))
            {
                using (StreamReader streamReader = new StreamReader(cryptoStream))
                {
                    ASCIIEncoding encoding = new ASCIIEncoding();
                    byte[] buffer = encoding.GetBytes(streamReader.ReadToEnd());
                    decryptedStream.Write(buffer, 0, buffer.Length);
                }
            }

            decryptedStream.Position = 0;
            return decryptedStream;
        }

        #endregion

        #region Private Members

        /// <summary>
        /// Gets the crypto transform.
        /// </summary>
        /// <param name="cryptType">Type of the crypt.</param>
        /// <returns></returns>
        private ICryptoTransform GetCryptoTransform(CryptType cryptType)
        {
            Rijndael rijndael = Rijndael.Create();
            this.InitializeAlgorithm(rijndael);

            if (cryptType == CryptType.Encrypt)
            {
                return rijndael.CreateEncryptor();
            }
            else if (cryptType == CryptType.Decrypt)
            {
                return rijndael.CreateDecryptor();
            }
            
            return null;            
        }

        /// <summary>
        /// Initializes the algorithm.
        /// </summary>
        /// <param name="rijndael">The rijndael.</param>
        private void InitializeAlgorithm(Rijndael rijndael)
        {
            SecurityKey securityKey = this.GetSecurityKey();
            rijndael.Key = ProtectedData.Unprotect(securityKey.Key, null, DataProtectionScope.CurrentUser);
            rijndael.IV = ProtectedData.Unprotect(securityKey.IV, null, DataProtectionScope.CurrentUser);
        }

        /// <summary>
        /// Gets the key.
        /// </summary>
        /// <returns>Returns SecurityKey</returns>
        private SecurityKey GetSecurityKey()
        {
            RegistryKey valKey = Registry.CurrentUser.OpenSubKey(this.registryKeyName);
            SecurityKey securityKey = new SecurityKey();
            if (valKey != null)
            {                
                securityKey.Key = valKey.GetValue(this.keyName) as byte[];
                securityKey.IV = valKey.GetValue(this.iVName) as byte[];
            }
            else
            {
                Rijndael rijndael = Rijndael.Create();
                rijndael.GenerateKey();
                rijndael.GenerateIV();
                securityKey.Key = ProtectedData.Protect(rijndael.Key, null,
                                                             DataProtectionScope.CurrentUser);
                securityKey.IV = ProtectedData.Protect(rijndael.IV, null,
                                                          DataProtectionScope.CurrentUser);
                this.StoreSecurityKey(securityKey);
            }

            return securityKey;
        }

        /// <summary>
        /// Stores the security key.
        /// </summary>
        /// <param name="securityKey">The security key.</param>
        private void StoreSecurityKey(SecurityKey securityKey)
        {
            string user = Environment.UserDomainName + "\\" + Environment.UserName;
            RegistrySecurity security = new RegistrySecurity();
            RegistryAccessRule rule = new RegistryAccessRule(user,
                    RegistryRights.CreateSubKey | RegistryRights.WriteKey | RegistryRights.ReadKey,
                    InheritanceFlags.ContainerInherit, PropagationFlags.None,
                    AccessControlType.Allow);
            security.AddAccessRule(rule);
            Registry.CurrentUser.CreateSubKey(this.registryKeyName, RegistryKeyPermissionCheck.ReadWriteSubTree, security);
            Registry.SetValue(this.registryKeyPath, this.keyName, securityKey.Key);
            Registry.SetValue(this.registryKeyPath, this.iVName, securityKey.IV);
        }

        #endregion 
    }
}
